package junior.db.dao.jdbc;

import com.google.common.collect.Lists;
import junior.db.dao.base.IEntityTypeAssigned;
import junior.db.dao.query.AbstractQuery;
import junior.db.dao.query.QuerySqlProvider;
import junior.db.dao.update.BatchInsertOption;
import junior.db.dao.update.BatchUpdateOption;
import junior.db.dao.update.BatchUpsertOption;
import junior.db.dao.update.DeleteOption;
import junior.db.dao.update.InsertOption;
import junior.db.dao.update.UpdateFieldOption;
import junior.db.dao.update.UpdateOption;
import junior.db.dao.update.UpdateSqlProvider;
import junior.db.dao.update.UpsertOption;
import junior.db.mapping.MappingProvider;
import junior.util.exception.ExceptionUtils;
import junior.util.reflect.ReflectUtils;
import junior.util.value.ValueUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.annotations.Param;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.dao.DataAccessException;
import org.springframework.dao.support.DataAccessUtils;
import org.springframework.jdbc.core.ColumnMapRowMapper;
import org.springframework.jdbc.core.ConnectionCallback;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapperResultSetExtractor;
import org.springframework.jdbc.support.GeneratedKeyHolder;
import org.springframework.jdbc.support.JdbcUtils;

import javax.sql.DataSource;
import java.beans.PropertyDescriptor;
import java.math.BigDecimal;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author wubo
 * @date 2024/9/3 14:01
 */
@Slf4j
public class JdbcDao implements InitializingBean {
    private static final int DEFAULT_BATCH_SIZE = 200;
    
    private int batchSize = DEFAULT_BATCH_SIZE;
    private String namespace;
    private DataSource dataSource;
    private JdbcTemplate template;
    private RowMapperFactory rowMapperFactory;
    
    public JdbcDao(DataSource dataSource) {
        this.dataSource = dataSource;
    }
    
    public void setNamespace(String namespace) {
        this.namespace = namespace;
    }
    
    public void setBatchSize(int batchSize) {
        this.batchSize = batchSize;
    }
    
    @Override
    public void afterPropertiesSet() throws Exception {
        template = new JdbcTemplate(dataSource);
        rowMapperFactory = new RowMapperFactory(MappingProvider.getTableMapping(namespace), MappingProvider.getColumnMapping(namespace));
    }
    
    /*--------- query methods ----------*/
    
    public <E, Q extends AbstractQuery & IEntityTypeAssigned> List<E> list(Q query) {
        query.setNamespace(namespace);
        
        List<Object> values = new ArrayList<>();
        String sql = QuerySqlProvider.query(query, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        return (List<E>) template.query(sql, rowMapperFactory.getRowMapper(query.getResultType().getType()), values.toArray());
    }
    
    public <E, Q extends AbstractQuery & IEntityTypeAssigned> List<E> listForObject(Q query, Class<E> resultType) {
        query.setNamespace(namespace);
        
        List<Object> values = new ArrayList<>();
        String sql = QuerySqlProvider.query(query, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        return (List<E>) template.query(sql, rowMapperFactory.getRowMapper(resultType), values.toArray());
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<Map<String, Object>> listForMap(Q query) {
        query.setNamespace(namespace);
        
        List<Object> values = new ArrayList<>();
        String sql = QuerySqlProvider.query(query, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        return template.queryForList(sql, values.toArray());
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<String> listForString(Q query) {
        return doListForType(query, String.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<Integer> listForInteger(Q query) {
        return doListForType(query, Integer.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<Long> listForLong(Q query) {
        return doListForType(query, Long.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<Boolean> listForBoolean(Q query) {
        return doListForType(query, Boolean.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<BigDecimal> listForBigDecimal(Q query) {
        return doListForType(query, BigDecimal.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<Date> listForDate(Q query) {
        return doListForType(query, Date.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<LocalDate> listForLocalDate(Q query) {
        return doListForType(query, LocalDate.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> List<LocalDateTime> listForLocalDateTime(Q query) {
        return doListForType(query, LocalDateTime.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> Long count(Q query) {
        if (!query.hasAnySelects()) {
            query.selectExpression("count(*)", null);
        }
        return doGetForType(query, Long.class);
    }
    
    public <E, Q extends AbstractQuery & IEntityTypeAssigned> E get(Q query) {
        List<E> list = list(query);
        return DataAccessUtils.singleResult(list);
    }
    
    public <E, Q extends AbstractQuery & IEntityTypeAssigned> E getForObject(Q query, Class<E> resultType) {
        List<E> list = listForObject(query, resultType);
        return DataAccessUtils.singleResult(list);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> Map<String, Object> getForMap(Q query) {
        List<Map<String, Object>> list = listForMap(query);
        return DataAccessUtils.singleResult(list);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> String getForString(Q query) {
        return doGetForType(query, String.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> Integer getForInteger(Q query) {
        return doGetForType(query, Integer.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> Long getForLong(Q query) {
        return doGetForType(query, Long.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> Boolean getForBoolean(Q query) {
        return doGetForType(query, Boolean.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> BigDecimal getForBigDecimal(Q query) {
        return doGetForType(query, BigDecimal.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> Date getForDate(Q query) {
        return doGetForType(query, Date.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> LocalDate getForLocalDate(Q query) {
        return doGetForType(query, LocalDate.class);
    }
    
    public <Q extends AbstractQuery & IEntityTypeAssigned> LocalDateTime getForLocalDateTime(Q query) {
        return doGetForType(query, LocalDateTime.class);
    }
    
    private <T, Q extends AbstractQuery & IEntityTypeAssigned> T doGetForType(Q query, Class<T> type) {
        List<T> list = doListForType(query, type);
        return DataAccessUtils.singleResult(list);
    }
    
    private <T, Q extends AbstractQuery & IEntityTypeAssigned> List<T> doListForType(Q query, Class<T> type) {
        query.setNamespace(namespace);
        
        List<Object> values = new ArrayList<>();
        String sql = QuerySqlProvider.query(query, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        return (List<T>) template.queryForList(sql, type, values.toArray());
    }
    
    /*--------- insert methods ----------*/
    
    /**
     * 批量插入
     *
     * @param option
     * @param entity
     * @return 影响记录数
     */
    public int insert(Object entity, InsertOption option) {
        option.setNamespace(namespace);
        List<Object> values = new ArrayList<>();
        String sql = UpdateSqlProvider.insert(entity, option, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        String[] generateKeyFields = option.getGenerateKeyFields();
        
        return doInsert(entity, sql, values, generateKeyFields);
    }
    
    /**
     * <h3>批量更新</h3>
     *
     * @param option
     * @param entities
     * @return 影响记录数
     */
    public <T> int[] batchInsert(List<T> entities, BatchInsertOption option) {
        option.setNamespace(namespace);
        
        List<List> values = new ArrayList<>();
        for (int i = 0; i < entities.size(); i++) {
            values.add(new ArrayList<>());
        }
        String sql = UpdateSqlProvider.batchInsert(entities, option, (value, index) -> values.get(index).add(value));
        
        return doBatchInsert(entities, sql, values, option.getGenerateKeyFields());
    }
    
    /**
     * <h3>插入或更新，支持设置自增主键</h3>
     *
     * @param option
     * @param entity
     * @return 影响记录数
     */
    public int upsert(@Param("entity") Object entity, @Param("option") UpsertOption option) {
        option.setNamespace(namespace);
        List<Object> values = new ArrayList<>();
        String sql = UpdateSqlProvider.upsert(entity, option, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        String[] generateKeyFields = option.getGenerateKeyFields();
        
        return doInsert(entity, sql, values, generateKeyFields);
    }
    
    /**
     * <h3>批量插入或更新, 支持设置自增主键</h3>
     *
     * @param option
     * @param entities
     * @return 影响记录数
     */
    public <T> int[] batchUpsert(@Param("entities") List<T> entities, @Param("option") BatchUpsertOption option) {
        option.setNamespace(namespace);
        
        List<List> values = new ArrayList<>();
        for (int i = 0; i < entities.size(); i++) {
            values.add(new ArrayList<>());
        }
        String sql = UpdateSqlProvider.batchUpsert(entities, option, (value, index) -> values.get(index).add(value));
        
        return doBatchInsert(entities, sql, values, option.getGenerateKeyFields());
    }
    
    private <T> int[] doBatchInsert(List<T> entities, String sql, List<List> values, String[] generateKeyFields) {
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            for (List list : values) {
                log.debug("===> params: {}", list);
            }
        }
        
        return template.execute(new ConnectionCallback<int[]>() {
            @Override
            public int[] doInConnection(Connection conn) throws SQLException, DataAccessException {
                List<Integer> rowsAffected = new ArrayList<>();
                
                boolean supportBatchUpdate = JdbcUtils.supportsBatchUpdates(conn);
                
                PreparedStatement ps = null;
                try {
                    ps = conn.prepareStatement(sql, generateKeyFields);
                    
                    int entityIndex = 0;
                    List<Object> subEntities = new ArrayList<>();
                    for (List<List> subList : Lists.partition(values, batchSize)) {
                        for (List sub : subList) {
                            for (int i = 0; i < sub.size(); i++) {
                                setValue(ps, sub.get(i), i + 1);
                            }
                            if (supportBatchUpdate) {
                                ps.addBatch();
                                subEntities.add(entities.get(entityIndex));
                            } else {
                                rowsAffected.add(ps.executeUpdate());
                                
                                if (generateKeyFields != null) {
                                    doGenerateKeys(Lists.newArrayList(entities.get(entityIndex)), generateKeyFields, ps);
                                }
                            }
                            entityIndex++;
                        }
                        if (supportBatchUpdate) {
                            int[] affectedRows = ps.executeBatch();
                            rowsAffected.addAll(Arrays.stream(affectedRows).boxed().collect(Collectors.toList()));
                            
                            if (generateKeyFields != null) {
                                doGenerateKeys(subEntities, generateKeyFields, ps);
                            }
                        }
                    }
                } catch (Exception e) {
                    throw ExceptionUtils.wrap2Runtime(e);
                } finally {
                    JdbcUtils.closeStatement(ps);
                }
                return rowsAffected.stream().mapToInt(i -> i.intValue()).toArray();
            }
        });
    }
    
    private int doInsert(Object entity, String sql, List<Object> values, String[] generateKeyFields) {
        if (generateKeyFields != null) {
            GeneratedKeyHolder keyHolder = new GeneratedKeyHolder();
            int affectedRows = template.update(psc -> {
                PreparedStatement ps = psc.prepareStatement(sql, generateKeyFields);
                for (int i = 0; i < values.size(); i++) {
                    setValue(ps, values.get(i), i + 1);
                }
                return ps;
            }, keyHolder);
            
            setEntityGenerateKeys(keyHolder.getKeyList(), entity, generateKeyFields);
            return affectedRows;
        } else {
            return template.update(sql, values.toArray());
        }
    }
    
    private int doBatchInsert(List entities, String sql, Connection conn, List<Object> values, String[] generateKeyFields) {
        PreparedStatement ps = null;
        try {
            ps = conn.prepareStatement(sql, generateKeyFields);
            for (int i = 0; i < values.size(); i++) {
                setValue(ps, values.get(i), i + 1);
            }
            int affectedRows = ps.executeUpdate();
            
            if (generateKeyFields != null) {
                doGenerateKeys(entities, generateKeyFields, ps);
            }
            
            return affectedRows;
        } catch (Exception e) {
            throw ExceptionUtils.wrap2Runtime(e);
        } finally {
            JdbcUtils.closeStatement(ps);
        }
    }
    
    private void doGenerateKeys(List entities, String[] generateKeyFields, PreparedStatement ps) throws SQLException {
        ResultSet keys = ps.getGeneratedKeys();
        if (keys != null) {
            try {
                RowMapperResultSetExtractor<Map<String, Object>> rse = new RowMapperResultSetExtractor(new ColumnMapRowMapper(), entities.size());
                List<Map<String, Object>> generatedKeys = rse.extractData(keys);
                setEntitiesGenerateKeys(generatedKeys, entities, generateKeyFields);
            } finally {
                JdbcUtils.closeResultSet(keys);
            }
        }
    }
    
    private void setEntityGenerateKeys(List<Map<String, Object>> keyList, Object entity, String[] generateKeyFields) {
        if (keyList.size() > 0) {
            Map<String, Object> keyMap = keyList.get(0);
            doSetEntityGenerateKeys(keyMap, entity, generateKeyFields);
        }
    }
    
    private void setEntitiesGenerateKeys(List<Map<String, Object>> keyList, List entities, String[] generateKeyFields) {
        if (keyList.size() > 0) {
            for (int i = 0; i < keyList.size(); i++) {
                Map<String, Object> keyMap = keyList.get(i);
                doSetEntityGenerateKeys(keyMap, entities.get(i), generateKeyFields);
            }
        }
    }
    
    private void doSetEntityGenerateKeys(Map<String, Object> keyMap, Object entities, String[] generateKeyFields) {
        if (keyMap != null) {
            int index = 0;
            for (Object value : keyMap.values()) {
                PropertyDescriptor pd = ReflectUtils.getPropertyDescriptor(entities.getClass(), generateKeyFields[index++]);
                setFieldValue(entities, pd, value);
            }
        }
    }
    
    /*--------- update methods ----------*/
    
    public int update(Object entity, UpdateOption option) {
        option.setNamespace(namespace);
        List<Object> values = new ArrayList<>();
        String sql = UpdateSqlProvider.update(entity, option, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        return template.update(sql, values.toArray());
    }
    
    public int updateField(Class<?> entityClass, UpdateFieldOption option) {
        option.setNamespace(namespace);
        List<Object> values = new ArrayList<>();
        String sql = UpdateSqlProvider.updateField(entityClass, option, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        return template.update(sql, values.toArray());
    }
    
    public <T> int[] batchUpdate(List<T> entities, BatchUpdateOption option) {
        option.setNamespace(namespace);
        
        Map<String, List<List>> sqlValuesMap = UpdateSqlProvider.batchUpdate(entities, option);
        
        return template.execute(new ConnectionCallback<int[]>() {
            @Override
            public int[] doInConnection(Connection conn) throws SQLException, DataAccessException {
                List<Integer> rowsAffected = new ArrayList<>();
                
                boolean supportBatchUpdate = JdbcUtils.supportsBatchUpdates(conn);
                
                for (Map.Entry<String, List<List>> entry : sqlValuesMap.entrySet()) {
                    if (log.isDebugEnabled()) {
                        log.debug("===> sql: {}", entry.getKey());
                    }
                    
                    PreparedStatement ps = conn.prepareStatement(entry.getKey());
                    try {
                        for (List<List> subList : Lists.partition(entry.getValue(), batchSize)) {
                            for (List<List> values : subList) {
                                if (log.isDebugEnabled()) {
                                    log.debug("===> params: {}", values);
                                }
                                for (int i=0; i<values.size(); i++) {
                                    setValue(ps, values.get(i), i+1);
                                }
                                if (supportBatchUpdate) {
                                    ps.addBatch();
                                } else {
                                    rowsAffected.add(ps.executeUpdate());
                                }
                            }
                            if (supportBatchUpdate) {
                                int[] affectedRows = ps.executeBatch();
                                rowsAffected.addAll(Arrays.stream(affectedRows).boxed().collect(Collectors.toList()));
                            }
                        }
                    } finally {
                        JdbcUtils.closeStatement(ps);
                    }
                }
                return rowsAffected.stream().mapToInt(i -> i.intValue()).toArray();
            }
        });
    }
    
    /*--------- delete methods ----------*/
    
    public int delete(Class<?> entityClass, DeleteOption option) {
        option.setNamespace(namespace);
        List<Object> values = new ArrayList<>();
        String sql = UpdateSqlProvider.delete(option, entityClass, value -> values.add(value));
        
        if (log.isDebugEnabled()) {
            log.debug("===> sql: {}", sql);
            log.debug("===> params: {}", values);
        }
        
        return template.update(sql, values.toArray());
    }
    
    private void setFieldValue(Object entity, PropertyDescriptor pd, Object value) {
        try {
            if (pd != null) {
                if (pd.getWriteMethod() != null) {
                    pd.getWriteMethod().invoke(entity, ValueUtils.convert(value, pd.getPropertyType()));
                } else {
                    pd.setValue(pd.getName(), ValueUtils.convert(value, pd.getPropertyType()));
                }
            }
        } catch (Exception e) {
            throw ExceptionUtils.wrap2Runtime(e);
        }
    }
    
    private void setValue(PreparedStatement ps, Object value, int index) throws SQLException {
        JdbcHelper.setValue(ps, value, index);
    }
    
}
